import os.path

import chromadb
from langchain_chroma import Chroma
from langchain_community.embeddings import ModelScopeEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from modelscope import HubApi
from uuid import uuid4


def load_kb_path(collection_name: str = "default"):
    current_path = os.path.abspath(__file__)
    project_root = os.path.dirname(os.path.dirname(current_path))
    return os.path.join(project_root, "knowledge_base", collection_name, "vector_store")


class KnowledgeBaseService:
    """provides access to the knowledge base"""
    def __init__(self):
        self.collection = None
        self.api = HubApi()
        self.api.login(os.getenv("MODEL_SCOPE_API_KEY"))
        self.embeddings = ModelScopeEmbeddings(
            model_id="iic/nlp_gte_sentence-embedding_chinese-large"
        )
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=200,
            chunk_overlap=50,
            length_function=len,
            separators=[
                "\n\n",
                "\n",
                " ",
                ".",
                ",",
                "。",  # Zero-width space
                "；",  # Fullwidth comma
                "\u3001",  # Ideographic comma
                "\uff0e",  # Fullwidth full stop
                "\u3002",  # Ideographic full stop
                "",
            ],
            # Existing args
        )

    def save_document(self, document: str, collection_name: str = "default"):
        """saves a document to a collection"""
        if self.collection is None or self.collection._collection_name != collection_name:
            self.collection = Chroma(
                collection_name=collection_name,
                embedding_function=self.embeddings,
                persist_directory="../knowledge_base/default/vector_store",
            )
        docs = self.text_splitter.create_documents([document])
        uuids = [str(uuid4()) for _ in range(len(docs))]
        return self.collection.add_documents(documents=docs, ids=uuids)

    def search_documents(self, query: str, collection_name: str = "default", top_k: int = 1):
        """searches for documents in a collection"""
        if self.collection is None or self.collection._collection_name != collection_name:
            self.collection = Chroma(
                collection_name=collection_name,
                embedding_function=self.embeddings,
                persist_directory=load_kb_path(),
            )
        documents = self.collection.similarity_search_with_score(query, k=top_k)
        return documents


knowledge_base_service = KnowledgeBaseService()


if __name__ == '__main__':
    kb_service = KnowledgeBaseService()
    # txt = """运动护眼六法：\n眨眨眼睛\n平时，在阅读或书写间隙，可用眨眨眼睛的办法来活动眼睛的保护肌。通过这样轻度上下磨擦眼皮还能增进眼睛的滋润。\n看看远处\n
    # 每天清晨，站在比较高的地方，两脚分开与肩同宽，双手叉腰，眼球左右移动，然后注视远处一个固定的目标，如树梢、旗杆、电视塔尖、云彩等。对这些远望的目标，要尽量看清楚，每次坚持一分钟。接着，再看近处一次，反复进行三次。长此以往，不仅可以
    # 消除眼肌疲劳，而且能够增强视觉功能。\n望望星空\n夏夜，空中常常是满天星斗，不妨抬头多多观赏，既可增长知识，又能使疲劳了一天的眼肌得到放松，真是一举两得！\n放放风筝\n到野外，迎着和煦的春风，经常放飞风筝，不仅能陶冶情操，锻炼身体，
    # 而且还能预防和治疗近视。\n打乒乓球\n要是经常打打乒乓，能让眼睛随乒乓球不停地进行远、近、上、下的调节和运动，促进眼球组织血液供应和代谢，有效地改善眼肌功能。\n做做“眼操”\n为了保护眼睛，应坚持每天认真做两次“眼保健操”。\n
    # 电脑者护眼小知识\n电脑在使人们的生活迈进信息化的同时，也给人类带来了新的疾病，被人们称之为“电脑眼病”。\n一项调查显示，美国一公司里的职员每天电脑前工作2小时以上，其中86％的人说在使用电脑时感到烦躁、疲劳、注意力难以集中、眼睛
    # 发干或是头疼，还有人出现了植物神经功能失调和忧郁症。美国全国职业保健与安全研究所的一项报告也表明每天在电脑前工作3小时以上者，90％的人眼睛有问题。\n电脑眼病的发病因素\n1.视觉疲劳：操作者眼睛在视屏、文件和键盘之间频繁移动，双
    # 眼不断地在各视点及视距间频敏调节，加上视屏的闪烁、反光和眩目造成视觉疲劳；\n2.精神紧张和神经失调：电脑作业是一项脑、眼、手密切配合的工作，具有紧张的脑力劳动和复杂的手工操作的共同特点，久而久之，导致精神高度紧张和神经失调；\n
    # 3.眼自洁能力减弱：眨眼次数只及平时的确1/3，因而减少了眼内润滑剂和酶的分泌。\n电脑可能导致的眼睛不适有哪些？\n眼睛疲劳\n如果室内的灯光或窗外的光线造成荧光幕反光时，便需要格外费神用力才能够看得清楚荧光幕上的字，在过度使用眼
    # 部肌肉的情况下，眼睛会很容易感觉疲劳。长时间集中精神对着电脑工作，眼部肌肉由于过度维持在同一个位置，亦会令眼睛疲倦。\n眼睛刺痛\n长时间近距离观看电脑荧光屏，会引致眼睛刺痛。一方面是由于减少了眨眼，眼睛缺乏泪水的滋润；另一方面，
    # 是因为受到电脑荧幕上不停闪动的小光点的刺激。\n眼睛干涩流眼泪\n眨眼会使泪水平均地分布在角膜表面，令角膜保持湿润，维持角膜的良好光学性质。而盯着电脑工作时，一般人都会比较少眨眼，因此眼睛便会缺乏泪水的滋润，从而眼睛干涩。眼睛干涩时，
    # 眼球会受到刺激，泪腺再分泌出来的泪水却流不到干的地方，反而会流出眼外。\n对焦缓慢\n长时间将视线专注于某一点，或是快速转动眼球，视线不停在文件、荧光屏、键盘上游移，都会使控制光线进入眼部的虹膜肌及主控对焦距离的睫状肌感到疲劳，因而
    # 造成眼睛不易对焦。若协调不足，还可能引发头痛。\n眼睛痒胀充血\n长时间盯着电脑工作时，眼球往往缺乏泪水的滋润，眼部肌肉也往往过度维持在同一个焦点或位置，会令眼睛痒胀充血。\n视觉模糊或重影\n不停对着电脑工作，眼睛疲劳、对焦困难，
    # 再加上在微弱的光线下对着电脑工作，易造成视线不能集中在目标物上，出现视觉模糊或重影。\n操作电脑的用眼之道\n1.每隔1小时休息一下双眼，伸展身体，不要持续操作。\n2.姿势和距离也是很重要的，尽量保持60厘米以上距离，
    # 调整一个最适当的姿势。\n3.调整屏幕的光度与清晰度，使之舒适悦眼；此外，屏幕过脏会使画面模糊，亦是造成眼睛疲劳的因素，所以要将屏幕拭擦干净。\n4.空调温度不要太低，出风口避免吹向头部。\n5.持续使用电脑时，要多眨眼睛以润湿双眼。\n
    # 6.屏幕反光会使眼睛疲劳，因此室内最好拉起窗帘，或是调整光线的投射角度。晚间使用电脑时，环境的光线以柔和为好。\n7.使用眼药水，消除眼睛疲劳。护眼小天使提醒你：如果你总觉得眼睛疲劳、颈肩酸痛，就要赶快注意一下你操作电脑的姿势了，
    # 也许换一个正确舒适的姿势，上网冲浪会更轻松。\n"""
    # kb_service.save_document(txt, "eye")
    results = kb_service.search_documents("使用电脑对眼睛的危害", "eye", 5)

    for res, score in results:
        print(f"* [SIM={score:3f}] {res.page_content} [{res.metadata}]")
